import cv2
import numpy as np



def q1():
    print("""1. 在测试视频(OpenCV安装目录\sources\samples\data)上，使用基于混合高斯模型的背景提取算法，
提取前景并显示(显示二值化图像，前景为白色)。""")

    #cap = cv2.VideoCapture('E:\python\data\vtest.avi')
    videoFileName = r'E:\python\data\vtest.avi'

    cap = cv2.VideoCapture(videoFileName)
    fgbg = cv2.createBackgroundSubtractorMOG2()
    while 1:
        ret, frame = cap.read()
        if not ret:
            break
        fgmask = fgbg .apply(frame)
        bg_img = fgbg .getBackgroundImage()

        cv2.imshow('mask', fgmask)
        cv2.imshow('bg img', bg_img)
        if cv2.waitKey(1) & 0xff == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()
    print(('\n'+'#'*40)*2+'\n')


def q2():
    print("2. 在1基础上，将前景目标进行分割，进一步使用不同颜色矩形框标记，并在命令行窗口中输出每个矩形框的位置和大小。")

    videoFileName = r'E:\python\data\vtest.avi'

    cap = cv2.VideoCapture(videoFileName)
    fgbg = cv2.createBackgroundSubtractorMOG2()
    thre = 200

    count = 0
    while 1:
        ret, frame = cap.read()
        if not ret:
            break
        fgmask  = fgbg .apply(frame)
        _, mask1 = cv2.threshold(fgmask, 30, 255, cv2.THRESH_BINARY)
        bg_img = fgbg .getBackgroundImage()
        _, cnts, _ = cv2.findContours(mask1, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if count > 0:
            print('\n' + '-' * 40, '\n这是第', count, '个frame')
            f1 = frame.copy()
            # cv2.drawContours(f1, cnts, -1, (0, 255, 0), 1)
            count1 = 0
            for i in cnts:
                area = cv2.contourArea(i)
                if area > thre:
                    x, y, w, h = cv2.boundingRect(i)
                    color = list(np.random.random(size=3) * 255)
                    cv2.rectangle(f1, (x, y), (x+w, y+h), color, 1)
                    print("第", count1, "个物体，")
                    print("矩形框的位置:", ((x, y), (x+w, y+h)), "矩形框的大小:", int(area))
                    print('-'*40)
                    count1 += 1
            # print(np.array(fgmask).shape)
            # exit()
            cv2.imshow('img', f1)
            cv2.imshow('mask', cv2.resize(fgmask, (384, 288)))
            cv2.imshow('bg img', bg_img)
        count += 1

        if (cv2.waitKey(2) == 27) or count > 10000:
            cap.release()
            cv2.destroyAllWindows()
            break
    print(('\n' + '#' * 40) * 2 + '\n')


def q4():
    print("4. 使用光流估计方法，在前述测试视频上计算特征点，进一步进行特征点光流估计。")

    videoFileName = r'E:\python\data\vtest.avi'

    cap = cv2.VideoCapture(videoFileName)

    # 角点检测所需参数
    feature_params = dict(maxCorners=100,
                          qualityLevel=0.3,
                          minDistance=7,
                          # corners=None, mask=None,
                          blockSize=7)
    # lucas kanade参数
    lk_params = dict(winSize=(15, 15),
                     maxLevel=2,
                     criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))

    _, frame = cap.read()
    prev_grey = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    # p0 = cv2.goodFeaturesToTrack(prev_grey, mask=None, **feature_params)
    p0 = cv2.goodFeaturesToTrack(prev_grey, **feature_params)

    while 1:
        ret, frame = cap.read()
        if not ret:
            break
        grey = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        p1, st, err = cv2.calcOpticalFlowPyrLK(prev_grey, grey, p0, None, **lk_params)
        # p1, st, err = cv2.calcOpticalFlowPyrLK(prev_grey,grey,p0,)

        goodPoints = p1[st == 1]    # length 56
        goodPrevPoints = p0[st == 1]
        result = frame.copy()
        color = (0, 0, 255)  # red
        for i in range(len(goodPoints)):
            curr, prev = goodPoints[i], goodPrevPoints[i]
            x0, y0 = curr.ravel()
            x1, y1 = prev.ravel()
            cv2.line(result, (x0, y0), (x1, y1), color)
            cv2.circle(result, (x0, y0), 3, color)

            prev_grey = grey.copy()
            p0 = goodPoints.reshape(-1, 1, 2)
            cv2.imshow('My result', result)

            if cv2.waitKey(30) == 27:
                cap.release()
                cv2.destroyAllWindows()
                break
    print(('\n' + '#' * 40) * 2 + '\n')


if __name__ == '__main__':
    q1()

    #q2()

    #q4()